package main

import (
	"bufio"
	"context"
	"crypto/rand"
	"fmt"
	"github.com/libp2p/go-libp2p"
	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/host"
	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"
	"github.com/libp2p/go-libp2p/core/protocol"
	"github.com/libp2p/go-libp2p/p2p/discovery/mdns"
	"github.com/libp2p/go-libp2p/p2p/host/autonat"
	"github.com/libp2p/go-libp2p/p2p/net/connmgr"
	"github.com/pion/stun"
	"io"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"
)

const protocolID = protocol.ID("/chat/1.0.0")
const discoveryNamespace = "libp2p-chat"

// 全局节点注册表 - 模拟简单的发现机制
var (
	globalPeers = make(map[peer.ID]peer.AddrInfo)
	peersMutex  sync.RWMutex
)

type discoveryNotifee struct {
	h host.Host
}

func (n *discoveryNotifee) HandlePeerFound(pi peer.AddrInfo) {
	if pi.ID == n.h.ID() {
		return
	}

	fmt.Printf("🔍 发现节点: %s\n", pi.ID.String())

	// 注册节点
	peersMutex.Lock()
	globalPeers[pi.ID] = pi
	peersMutex.Unlock()

	// 异步连接，让hole punching发挥作用
	go func() {
		// 稍微延迟，让双方都有时间准备
		time.Sleep(time.Duration(100+len(pi.ID)%500) * time.Millisecond)

		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
		defer cancel()

		if err := n.h.Connect(ctx, pi); err != nil {
			fmt.Printf("❌ 连接失败: %s\n", err)
		} else {
			fmt.Printf("✅ 成功连接到: %s\n", pi.ID.String())
		}
	}()
}

func main() {

	getMappedAddr("stun.l.google.com:19302")
	getMappedAddr("stun1.l.google.com:19302")
	//if len(os.Args) < 2 {
	//	fmt.Println("libp2p NAT穿透聊天工具")
	//	fmt.Println()
	//	fmt.Println("用法:")
	//	fmt.Println("  go run main.go server           # 启动服务器")
	//	fmt.Println("  go run main.go client <peer-id> # 连接到指定节点")
	//	fmt.Println("  go run main.go list             # 列出发现的节点")
	//	return
	//}

	//mode := os.Args[1]
	mode := "server"
	// 创建支持NAT穿透的libp2p主机
	h, err := createNATHost()
	if err != nil {
		log.Fatal("创建主机失败:", err)
	}
	defer h.Close()

	fmt.Printf("🚀 节点启动成功\n")
	fmt.Printf("📍 节点ID: %s\n", h.ID().String())

	// 设置协议处理器
	h.SetStreamHandler(protocolID, handleStream)

	// 启动发现服务
	startDiscovery(h)

	ctx := context.Background()

	switch mode {
	case "server":
		fmt.Printf("🎯 服务器模式启动\n")
		fmt.Printf("💡 其他节点可以通过以下命令连接:\n")
		fmt.Printf("   go run main.go client %s\n", h.ID().String())
		fmt.Printf("💬 可以直接输入消息发送给连接的客户端\n")
		fmt.Println()

		// 服务器也可以发送消息
		go handleUserInput(h)

		select {} // 保持运行

	case "client":
		if len(os.Args) < 3 {
			fmt.Println("❌ 需要指定目标节点ID")
			return
		}

		targetPeerID := os.Args[2]
		fmt.Printf("🎯 客户端模式，目标: %s\n", targetPeerID)

		if err := connectToTarget(ctx, h, targetPeerID); err != nil {
			log.Fatal("连接失败:", err)
		}

	case "list":
		listPeers(h)

	default:
		fmt.Println("❌ 无效模式")
	}
}

func getMappedAddr(server string) {
	conn, _ := net.Dial("udp", server)
	c, _ := stun.NewClient(conn)
	defer c.Close()

	var xorAddr stun.XORMappedAddress
	_ = c.Do(stun.MustBuild(stun.TransactionID, stun.BindingRequest),
		func(res stun.Event) {
			_ = xorAddr.GetFrom(res.Message)
			fmt.Printf("STUN server %s 映射地址: %s\n", server, xorAddr.String())
		},
	)
}

// 创建支持NAT穿透的主机
func createNATHost() (host.Host, error) {
	priv, _, err := crypto.GenerateKeyPairWithReader(crypto.RSA, 2048, rand.Reader)
	if err != nil {
		return nil, err
	}

	connmgr, err := connmgr.NewConnManager(5, 50)
	if err != nil {
		return nil, err
	}

	// 关键：启用所有NAT穿透功能
	opts := []libp2p.Option{
		libp2p.Identity(priv),
		// 监听所有接口
		libp2p.ListenAddrStrings(
			"/ip4/0.0.0.0/tcp/0",
		),
		libp2p.ConnectionManager(connmgr),
		// 核心NAT穿透功能
		libp2p.EnableNATService(),   // NAT服务
		libp2p.EnableHolePunching(), // Hole punching
		libp2p.NATPortMap(),         // UPnP端口映射

		// 默认协议栈
		libp2p.DefaultSecurity,
		libp2p.DefaultMuxers,
		libp2p.DefaultTransports,
	}

	h, err := libp2p.New(opts...)
	if err != nil {
		return nil, err
	}
	c, err := autonat.New(h)
	fmt.Println(c.Status())
	fmt.Printf("🌐 NAT穿透功能已启用\n")
	fmt.Printf("📡 监听地址: %v\n", h.Addrs())

	return h, nil
}

// 启动发现服务
func startDiscovery(h host.Host) {
	// mDNS本地发现
	disc := mdns.NewMdnsService(h, discoveryNamespace, &discoveryNotifee{h: h})
	if err := disc.Start(); err != nil {
		fmt.Printf("⚠️ mDNS启动失败: %s\n", err)
	} else {
		fmt.Printf("🔍 本地发现服务已启动\n")
	}

	// 定期显示连接状态
	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()

		for range ticker.C {
			peers := h.Network().Peers()
			if len(peers) > 0 {
				fmt.Printf("📊 当前连接: %d 个节点\n", len(peers))
			}
		}
	}()
}

// 处理流连接
func handleStream(s network.Stream) {
	remotePeer := s.Conn().RemotePeer().String()
	fmt.Printf("📞 新连接来自: %s\n", remotePeer)

	// 显示连接路径（用于调试NAT穿透）
	conn := s.Conn()
	fmt.Printf("🔗 连接路径: %s -> %s\n",
		conn.LocalMultiaddr(), conn.RemoteMultiaddr())

	reader := bufio.NewReader(s)

	go func() {
		defer s.Close()
		for {
			msg, err := reader.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					fmt.Printf("❌ 读取错误: %s\n", err)
				}
				return
			}
			msg = strings.TrimSpace(msg)
			if msg != "" {
				fmt.Printf("📨 [%s]: %s\n", remotePeer[:8], msg)
			}
		}
	}()
}

// 处理用户输入（服务器模式）
func handleUserInput(h host.Host) {
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		input := scanner.Text()
		if input == "" {
			continue
		}
		if input == "quit" {
			os.Exit(0)
		}
		if input == "list" {
			listPeers(h)
			continue
		}

		// 广播消息给所有连接的节点
		peers := h.Network().Peers()
		if len(peers) == 0 {
			fmt.Printf("⚠️ 没有连接的节点\n")
			continue
		}

		for _, peerID := range peers {
			go func(pid peer.ID) {
				stream, err := h.NewStream(context.Background(), pid, protocolID)
				if err != nil {
					fmt.Printf("❌ 无法向 %s 发送消息: %s\n", pid.String()[:8], err)
					return
				}
				defer stream.Close()

				writer := bufio.NewWriter(stream)
				if _, err := writer.WriteString(input + "\n"); err != nil {
					fmt.Printf("❌ 发送失败: %s\n", err)
					return
				}
				writer.Flush()
			}(peerID)
		}
		fmt.Printf("📤 消息已发送给 %d 个节点\n", len(peers))
	}
}

// 连接到目标节点
func connectToTarget(ctx context.Context, h host.Host, targetPeerID string) error {
	targetID, err := peer.Decode(targetPeerID)
	if err != nil {
		return fmt.Errorf("无效的节点ID: %w", err)
	}

	fmt.Printf("🔍 搜索目标节点...\n")

	// 等待发现目标节点
	found := false
	for i := 0; i < 30; i++ { // 等待30秒
		peersMutex.RLock()
		if _, exists := globalPeers[targetID]; exists {
			found = true
			peersMutex.RUnlock()
			break
		}
		peersMutex.RUnlock()

		// 检查是否已经连接
		if h.Network().Connectedness(targetID) == network.Connected {
			found = true
			break
		}

		fmt.Print(".")
		time.Sleep(1 * time.Second)
	}
	fmt.Println()

	if !found {
		return fmt.Errorf("未找到目标节点。请确保：\n1. 目标节点正在运行\n2. 双方在同一网络或NAT穿透成功")
	}

	// 等待连接建立
	for i := 0; i < 10; i++ {
		if h.Network().Connectedness(targetID) == network.Connected {
			break
		}
		time.Sleep(500 * time.Millisecond)
	}

	if h.Network().Connectedness(targetID) != network.Connected {
		return fmt.Errorf("无法连接到目标节点")
	}

	fmt.Printf("✅ 已连接到目标节点\n")

	// 显示连接信息
	conns := h.Network().ConnsToPeer(targetID)
	for _, conn := range conns {
		fmt.Printf("🔗 连接路径: %s -> %s\n",
			conn.LocalMultiaddr(), conn.RemoteMultiaddr())
	}

	// 开始聊天
	return startChat(ctx, h, targetID)
}

// 开始聊天会话
func startChat(ctx context.Context, h host.Host, targetID peer.ID) error {
	stream, err := h.NewStream(ctx, targetID, protocolID)
	if err != nil {
		return fmt.Errorf("创建流失败: %w", err)
	}
	defer stream.Close()

	fmt.Printf("💬 聊天开始！输入消息按回车发送，输入 'quit' 退出\n")
	fmt.Println()

	// 接收消息
	go func() {
		reader := bufio.NewReader(stream)
		for {
			msg, err := reader.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					fmt.Printf("❌ 接收错误: %s\n", err)
				}
				return
			}
			msg = strings.TrimSpace(msg)
			if msg != "" {
				fmt.Printf("📨 对方: %s\n", msg)
			}
		}
	}()

	// 发送消息
	writer := bufio.NewWriter(stream)
	scanner := bufio.NewScanner(os.Stdin)

	for scanner.Scan() {
		input := scanner.Text()
		if input == "quit" {
			break
		}
		if input == "" {
			continue
		}

		if _, err := writer.WriteString(input + "\n"); err != nil {
			fmt.Printf("❌ 发送失败: %s\n", err)
			break
		}
		if err := writer.Flush(); err != nil {
			fmt.Printf("❌ 刷新失败: %s\n", err)
			break
		}
	}

	return scanner.Err()
}

// 列出发现的节点
func listPeers(h host.Host) {
	fmt.Printf("📋 节点列表:\n")

	peersMutex.RLock()
	defer peersMutex.RUnlock()

	if len(globalPeers) == 0 {
		fmt.Printf("   (暂无发现的节点)\n")
		return
	}

	for id, info := range globalPeers {
		connected := h.Network().Connectedness(id) == network.Connected
		status := "❌"
		if connected {
			status = "✅"
		}
		fmt.Printf("   %s %s (地址数: %d)\n", status, id.String(), len(info.Addrs))
	}
}
